package com.luncheng.mcp.core;

import com.luncheng.mcp.core.factory.McpServerFactory;
import com.luncheng.mcp.core.load.McpConfigLoader;
import com.luncheng.mcp.core.model.McpServerConfig;
import com.luncheng.mcp.core.router.DynamicRouterProvider;
import io.modelcontextprotocol.server.McpSyncServer;
import io.modelcontextprotocol.server.transport.WebMvcSseServerTransportProvider;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.DependsOn;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class McpSseServerBoostrap {
    @Autowired()
    private McpConfigLoader configLoader;
    @Autowired
    private DynamicRouterProvider dynamicRouterProvider;

    private static final Map<String, McpSyncServer> serverMap = new ConcurrentHashMap<>();
    private static final Map<String, WebMvcSseServerTransportProvider> providerMap = new ConcurrentHashMap<>();

    @PostConstruct
    public void start() {
        try {
            log.info("正在初始化 MCP 服务器...");
            if (Objects.nonNull(configLoader)) {
                for (McpServerConfig config : configLoader.loadServerConfig()) {
                    addMcpSyncServer(config);
                }
            }
            log.info("MCP 服务器初始化完成，共加载 {} 个服务", serverMap.size());
        } catch (Exception e) {
            log.error("MCP 服务器初始化失败", e);
            throw new RuntimeException("MCP 服务器初始化失败", e);
        }
    }

    public void addMcpSyncServer(McpServerConfig mcpServerConfig) {
        log.info("正在添加 MCP 服务器: {}", mcpServerConfig.getName());
        try {
            McpSyncServer mcpSyncServer = McpServerFactory.createServer(mcpServerConfig, transportProvider -> {
                providerMap.put(mcpServerConfig.getName(), transportProvider);
            });
            serverMap.put(mcpServerConfig.getName(), mcpSyncServer);
            log.info("MCP 服务器 [{}] 添加成功", mcpServerConfig.getName());

            // 刷新路由
            dynamicRouterProvider.refreshRouter();
        } catch (Exception e) {
            log.error("创建 MCP 服务器 [{}] 失败", mcpServerConfig.getName(), e);
            throw new RuntimeException("创建 MCP 服务器失败: " + mcpServerConfig.getName(), e);
        }
    }

    public void removeMcpSyncServer(String name) {
        log.info("正在移除 MCP 服务器: {}", name);
        McpSyncServer mcpSyncServer = serverMap.remove(name);
        if (mcpSyncServer != null) {
            WebMvcSseServerTransportProvider provider = providerMap.remove(name);
            mcpSyncServer.close();
            log.info("MCP 服务器 [{}] 已关闭并移除", name);

            // 刷新路由
            dynamicRouterProvider.refreshRouter();
        } else {
            log.warn("尝试移除不存在的 MCP 服务器: {}", name);
        }
    }

    @PreDestroy
    public void stop() {
        log.info("正在停止所有 MCP 服务器...");
        serverMap.forEach((name, server) -> {
            log.info("正在关闭 MCP 服务器: {}", name);
            server.close();
        });
        serverMap.clear();
        providerMap.clear();
        log.info("所有 MCP 服务器已停止");
    }

    public static Map<String, McpSyncServer> getServerMap() {
        return serverMap;
    }

    public static Map<String, WebMvcSseServerTransportProvider> getProviderMap() {
        return providerMap;
    }
}